#pragma once
#include <utility>
#include <assert.h>
using std::pair;
using std::make_pair;

enum Color{
  RED,
  BLACK
};

template <class T>
struct RBTreeNode {
  typedef RBTreeNode<T> Node;
  Node *_left;
  Node *_right;
  Node *_parent;
  T _data;
  Color _color;

  RBTreeNode(const T &data = T(), Color color = RED)
    :_left(nullptr),
    _right(nullptr),
    _parent(nullptr),
    _data(data),
    _color(color)
  {}
};

template <class T, class Ref, class Ptr>
class RBT_iterator{
  typedef RBTreeNode<T> Node;
  typedef RBT_iterator<T, Ref, Ptr> iterator;
  Node *_pnode;

public:
  RBT_iterator(Node *pnode)
   :_pnode(pnode)
  {}
  
  Ref operator*() const{
    return _pnode->_data;
  }

  Ptr operator->() const{
    return &_pnode->_data;
  }

  bool operator==(const iterator &it) const{
    return _pnode == it._pnode;
  }

  bool operator!=(const iterator &it) const{
    return _pnode != it._pnode;
  }

  iterator& operator++(){
    if(_pnode->_right != nullptr)
    {
      Node *left = _pnode->_right; 
      while(left->_left != nullptr)
      {
        left = left->_left;
      }
      _pnode = left;
    }
    else
    {
      Node *parent = _pnode->_parent;
      while(_pnode == parent->_right)
      {
        _pnode = parent;
        parent = parent->_parent;
      }
       
      if(_pnode->_right != parent)
        _pnode = parent;
    }
    return *this;
  }
   
  iterator operator++(int){
    iterator it(_pnode);
    ++*this;
    return it;
  }

  iterator operator--(){
    if(_pnode->_parent->_parent == _pnode && _pnode->_color == RED)
    {
      _pnode = _pnode->_right;
    }
    else if(_pnode->_left != nullptr)
    {
      Node *right = _pnode->_left;
      while(right->_right != nullptr)
      {
        right = right->_right;
      }
      _pnode = right;
    }
    else
    {
      Node *parent = _pnode->_parent;
      while(_pnode == parent->_left)
      {
        _pnode = parent;
        parent = parent->_parent;
      }
      _pnode = parent;
    }
    return *this;
  }

  iterator operator--(int){
    iterator it(_pnode);
    --*this;
    return it;
  }
  
};

template <class K, class T, class KofT>
class RBTree{
  typedef RBTreeNode<T> Node;
  Node *_phead;

public:
  typedef RBT_iterator<T, T&, T*> iterator;
  typedef RBT_iterator<T, const T&, const T*> const_iterator;

  RBTree(){
    _phead = new Node(T(), RED);
    _phead->_left = _phead->_right = _phead;
  }

  iterator begin(){
    return iterator(_phead->_left); 
  }
   
  iterator end(){
    return iterator(_phead);
  }

  const_iterator begin() const{
    return const_iterator(_phead->_left); 
  }
   
  const_iterator end() const{
    return const_iterator(_phead);
  }

  std::pair<iterator, bool> Insert(const T& data){
    Node* &rproot = GetRoot(); 
    if(rproot == nullptr)
    {
      rproot = new Node(data, BLACK);
      rproot->_parent = _phead;
      _phead->_left = _phead->_right = rproot;
      return make_pair(iterator(rproot), true);
    }

    KofT kot;
    Node *cur = rproot;
    Node *parent = nullptr;
    while(cur != nullptr)
    {
      parent = cur;
      if(kot(data) > kot(cur->_data))
        cur = cur->_right;
      else if(kot(data) < kot(cur->_data))
        cur = cur->_left;
      else
        return make_pair(iterator(cur), false);
    }

    cur = new Node(data, RED);
    if(kot(data) > kot(parent->_data))
    {
      parent->_right = cur;
    }
    else
    {
      parent->_left = cur;
    }
    cur->_parent = parent;

    Node* newnode = cur;
    while(parent != _phead && parent->_color == RED)
    {
      Node* grandparent = parent->_parent;
      assert(grandparent!=nullptr);
      assert(grandparent->_color == BLACK);

      Node* uncle = grandparent->_left;
      if(grandparent->_left == parent)
        uncle = grandparent->_right;

      if(uncle != nullptr && uncle->_color == RED)
      {
        parent->_color = uncle->_color = BLACK;
        grandparent->_color = RED;
        cur = grandparent;
        parent = grandparent->_parent;
      }
      else
      {
        if(parent == grandparent->_left)
        {
          if(cur == parent->_left)
          {
            RotateR(grandparent);
            parent->_color = BLACK;
            grandparent->_color = RED;
          }
          else
          {
            RotateL(parent);
            RotateR(grandparent);
            cur->_color = BLACK;
            grandparent->_color = RED;
          }
        }
        else
        {
          if(cur == parent->_right)
          {
            RotateL(grandparent);
            parent->_color = BLACK;
            grandparent->_color = RED;
          }
          else
          {
            RotateR(parent);
            RotateL(grandparent);
            cur->_color = BLACK;
            grandparent->_color = RED;
          }
        }
        break;
      }
    } //end of while

    if(parent == _phead)
      cur->_color = BLACK;
    _phead->_left = LeftMost();
    _phead->_right = RightMost();
    return make_pair(iterator(newnode), true);
  }

  iterator Find(const K &k)
  {
    KofT kot;
    Node *cur = GetRoot();
    if(cur == nullptr) return end();
    while(cur != nullptr)
    {
      if(k > kot(cur->_data))
        cur = cur->_right;
      else if(k < kot(cur->_data))
        cur = cur->_left;
      else
        return iterator(cur);
    }
    return end();
  }
  
private:
  Node*& GetRoot(){
    return _phead->_parent;
  }

  Node* LeftMost(){
    Node *proot = GetRoot();
    if(proot == nullptr)
    {
      return _phead;
    }
    else{
      Node *left = proot;
      while(left->_left != nullptr)
      {
        left = left->_left;
      }
      return left;
    }
  }

  Node* RightMost(){
    Node *proot = GetRoot();
    if(proot == nullptr)
    {
      return _phead;
    }
    else{
      Node *right = proot;
      while(right->_right != nullptr)
      {
        right = right->_right;
      }
      return right;
    }
  }

  void RotateL(Node *parent){
    Node* subR = parent->_right;
    Node* subRL = subR->_left;
    Node* ppNode = parent->_parent;

    subR->_left = parent;
    parent->_parent = subR;

    parent->_right = subRL;
    if(subRL != nullptr)
      subRL->_parent = parent;

    if(ppNode == _phead)
    {
      _phead->_parent = subR;
    }
    else
    {
      if(parent == ppNode->_left)
        ppNode->_left = subR; 
      else
        ppNode->_right = subR;
    }
    subR->_parent = ppNode;
  }
   
  void RotateR(Node *parent){
    Node *subL = parent->_left;
    Node *subLR = subL->_right;
    Node *ppNode = parent->_parent;

    subL->_right = parent;
    parent->_parent = subL;

    parent->_left = subLR;
    if(subLR != nullptr)
      subLR->_parent = parent;

    if(ppNode == _phead)
    {
      _phead->_parent = subL;
    }
    else
    {
      if(parent == ppNode->_left)
        ppNode->_left = subL;
      else
        ppNode->_right = subL;
    }
    subL->_parent = ppNode;
  }
};

